import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import argparse
from pathlib import Path
from scipy.signal import find_peaks
import json

def generate_jump_rope_visualization(csv_path, json_path, output_path):
    """
    根据跳绳运动的CSV数据和JSON检测结果，生成一个详细的分析可视化报告。
    该版本特点：
    1. 使用JSON文件中的'jump_frames'作为权威的跳跃计数和位置。
    2. 采用智能文本标注，通过上下交错的方式防止计数的数字重叠。
    3. 图表宽度根据视频时长动态调整，避免内容拥挤。
    4. 合并视图采用归一化坐标，方便对比不同曲线。
    5. 增加保护序列区间标注功能。
    """
    csv_path = Path(csv_path)
    json_path = Path(json_path)
    if not csv_path.exists():
        print(f"错误: CSV文件未找到 at {csv_path}")
        return
    if not json_path.exists():
        print(f"错误: JSON文件未找到 at {json_path}")
        return

    df = pd.read_csv(csv_path)
    with open(json_path, 'r', encoding='utf-8') as f:
        detection_result = json.load(f)

    valid_frames = df[df['has_pose'] == True].copy()

    if len(valid_frames) < 50:
        print(f"警告: {csv_path.name} 的有效姿态数据过少，无法生成图表。")
        return

    # --- 1. 数据准备 ---
    frame_indices = valid_frames.index.values
    com_y_raw = valid_frames['center_of_mass_y'].values
    ankle_y_raw = valid_frames['avg_ankle_y'].values
    wrist_y_raw = valid_frames['avg_wrist_y'].values

    # 1.1 数据归一化 (用于合并视图)
    def min_max_scale(data):
        min_val = np.min(data)
        max_val = np.max(data)
        if max_val - min_val == 0:
            return np.zeros_like(data)
        return (data - min_val) / (max_val - min_val)

    com_y_scaled = min_max_scale(com_y_raw)
    ankle_y_scaled = min_max_scale(ankle_y_raw)
    wrist_y_scaled = min_max_scale(wrist_y_raw)

    # --- 2. 核心算法：从JSON文件加载权威的跳跃峰值 ---
    jump_frames_data = detection_result.get('jump_frames', [])
    if not jump_frames_data:
        print("警告: JSON文件中没有找到 'jump_frames' 数据，无法标注跳跃。")
        peaks = []
    else:
        # 从 'jump_frames' 提取帧号
        peaks = [item['frame'] for item in jump_frames_data]
    
    authoritative_count = len(peaks)

    # --- 3. 图表绘制 ---
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 3.1 动态计算图表宽度和高度
    # 密度：每30帧=2cm。 1英寸=2.54cm。
    num_valid_frames = len(valid_frames)
    width_cm = (num_valid_frames / 30.0) * 2 # 核心修改：每30帧的长度从1cm增加到2cm
    width_inches = width_cm / 2.54
    # 设置一个合理的最小宽度，移除宽度上限
    width_inches = max(20, width_inches) # 最小20英寸，无最大宽度限制
    
    # 动态计算高度，保持合理的宽高比
    # 基础高度为16英寸，根据宽度适当调整
    base_height = 16
    if width_inches > 40:
        # 对于很宽的图表，适当增加高度
        height_inches = base_height + (width_inches - 40) * 0.2
        height_inches = min(height_inches, 32)  # 最大高度限制为32英寸
    else:
        height_inches = base_height
    
    fig, axes = plt.subplots(4, 1, figsize=(width_inches, height_inches), sharex=True)
    
    # 全局标题
    video_name = csv_path.stem.replace('_frame_data', '')
    main_title = f'跳绳综合分析: {video_name} - 权威计数: {authoritative_count} 次'
    fig.suptitle(main_title, fontsize=20, fontweight='bold')

    # --- 3.2 添加保护序列区间标注 ---
    protected_sequences = detection_result.get('protected_sequences', [])
    for i, seq in enumerate(protected_sequences):
        start_frame = seq['start_frame']
        end_frame = seq['end_frame']
        confidence = seq.get('confidence', 0.0)
        
        # 为每个子图添加保护序列背景色块
        for ax in axes:
            ax.axvspan(start_frame, end_frame, alpha=0.2, color='lightgreen', 
                      label=f'保护序列 {i+1}' if ax == axes[0] else "")

    # 子图1: 合并数据视图 (相对坐标系)
    ax1 = axes[0]
    ax1.plot(frame_indices, com_y_scaled, color='#1f77b4', label='重心 (相对)')
    ax1.plot(frame_indices, ankle_y_scaled, color='#2ca02c', label='脚踝 (相对)')
    ax1.plot(frame_indices, wrist_y_scaled, color='#ff7f0e', label='手腕 (相对)')
    ax1.set_title(f'合并数据视图 (相对坐标系)', fontsize=14)
    ax1.set_ylabel('相对位置 (归一化)', fontsize=12)
    ax1.legend(loc='upper right')
    ax1.grid(True, alpha=0.6)

    # 在第一个子图上添加保护序列文本标注
    for i, seq in enumerate(protected_sequences):
        start_frame = seq['start_frame']
        end_frame = seq['end_frame']
        confidence = seq.get('confidence', 0.0)
        mid_frame = (start_frame + end_frame) / 2
        ax1.text(mid_frame, 0.9, f'保护序列{i+1}\n置信度:{confidence:.3f}', 
                ha='center', va='top', fontsize=10, 
                bbox=dict(boxstyle="round,pad=0.3", facecolor='lightgreen', alpha=0.7))

    # 子图2: 重心 (Y轴)
    ax2 = axes[1]
    ax2.plot(frame_indices, com_y_raw, color='#ff7f0e', label='原始重心Y')
    ax2.set_title(f'重心 (Y轴)', fontsize=14)
    ax2.set_ylabel('位置', fontsize=12)
    ax2.legend(loc='upper right')
    ax2.grid(True, alpha=0.6)

    # 子图3: 平均脚踝位置 (Y轴)
    ax3 = axes[2]
    ax3.plot(frame_indices, ankle_y_raw, color='#2ca02c', label='原始平均脚踝Y')
    ax3.set_title(f'平均脚踝位置 (Y轴)', fontsize=14)
    ax3.set_ylabel('位置', fontsize=12)
    ax3.legend(loc='upper right')
    ax3.grid(True, alpha=0.6)

    # 子图4: 平均手腕位置 (Y轴)
    ax4 = axes[3]
    ax4.plot(frame_indices, wrist_y_raw, color='#9400d3', label='原始平均手腕Y')
    ax4.set_title(f'平均手腕位置 (Y轴)', fontsize=14)
    ax4.set_ylabel('位置', fontsize=12)
    ax4.legend(loc='upper right')
    ax4.grid(True, alpha=0.6)
    
    # --- 4. 智能标注：实现数字不重叠 ---
    # 定义一个垂直偏移量，用于将数字上下错开
    y_offset_stagger = 0.05 # 可以调整这个值来改变数字的偏移距离

    for i, frame in enumerate(peaks):
        # 找到原始dataframe中与frame号最接近的索引
        peak_idx = np.abs(frame_indices - frame).argmin()

        # 根据计数的奇偶性决定数字是在点的上方还是下方
        if i % 2 == 0:
            # 偶数，数字在上方
            vertical_offset = y_offset_stagger
            va = 'bottom'  # 垂直对齐方式
        else:
            # 奇数，数字在下方
            vertical_offset = -y_offset_stagger
            va = 'top'
        
        # 标注1: 合并图 (以重心为基准, 使用归一化数据)
        com_y_scaled_val = com_y_scaled[peak_idx]
        ax1.plot(frame, com_y_scaled_val, 'rx', markersize=8, markeredgewidth=1.5)
        ax1.text(frame, com_y_scaled_val + vertical_offset, str(i + 1), color='red', fontsize=10, ha='center', va=va)

        # 标注2: 重心图 (使用原始数据)
        com_y_val = com_y_raw[peak_idx]
        ax2.plot(frame, com_y_val, 'rx', markersize=8, markeredgewidth=1.5)
        ax2.text(frame, com_y_val + vertical_offset, str(i + 1), color='red', fontsize=10, ha='center', va=va)

        # 标注3: 脚踝图 (使用原始数据)
        ankle_y_val = ankle_y_raw[peak_idx]
        ax3.plot(frame, ankle_y_val, 'rx', markersize=8, markeredgewidth=1.5)
        ax3.text(frame, ankle_y_val + vertical_offset, str(i + 1), color='red', fontsize=10, ha='center', va=va)

        # 标注4: 手腕图 (使用原始数据)
        wrist_y_val = wrist_y_raw[peak_idx]
        ax4.plot(frame, wrist_y_val, 'rx', markersize=8, markeredgewidth=1.5)
        ax4.text(frame, wrist_y_val + vertical_offset, str(i + 1), color='red', fontsize=10, ha='center', va=va)

    # 翻转所有Y轴
    for ax in axes:
        ax.invert_yaxis()

    plt.xlabel('帧号', fontsize=12)
    plt.tight_layout(rect=[0, 0.03, 1, 0.95], h_pad=3.0) # 增加子图间距和调整边距
    
    plt.savefig(output_path, dpi=300, bbox_inches='tight', facecolor='white', pad_inches=0.5)
    plt.close()
    print(f"分析图表已生成并保存到: {output_path}")

def main():
    parser = argparse.ArgumentParser(description='为跳绳数据生成一个基于JSON的、带防重叠标注的分析可视化报告。')
    parser.add_argument('--csv', required=True, help='包含逐帧关键点数据的CSV文件路径。')
    parser.add_argument('--json', required=True, help='包含权威检测结果的JSON文件路径。')
    parser.add_argument('--output', required=True, help='生成的图表图片的保存路径。')
    args = parser.parse_args()

    generate_jump_rope_visualization(args.csv, args.json, args.output)

if __name__ == "__main__":
    main()